#!/usr/bin/env python3
"""
Phase 3C — Scene Engine 镜头引擎验证
Playwright 帧序列渲染器

用途:
  使用 Playwright 打开 HTML/CSS 动画模板，按 15fps 截图，输出 PNG 帧序列。
  然后由 FFmpeg 合成为 H.264 mp4。

用法:
  python render_phase3c_scene_engine.py

输出:
  - 03_rendered_frames/{scene_name}/frame_XXXXX.png
  - 04_scene_clips/{scene_name}.mp4
  - 05_contact_sheets/{scene_name}_contact.jpg
"""

import os
import sys
import time
import subprocess

# ── 路径配置 ──
PROJECT_ROOT = r"E:\集群文件夹\factory_os\short_video_real_data_pipeline\phase3c_scene_engine_test"
HTML_DIR     = os.path.join(PROJECT_ROOT, "02_html_animation")
FRAMES_DIR   = os.path.join(PROJECT_ROOT, "03_rendered_frames")
CLIPS_DIR    = os.path.join(PROJECT_ROOT, "04_scene_clips")
CONTACTS_DIR = os.path.join(PROJECT_ROOT, "05_contact_sheets")

# ── Scene 定义 ──
SCENES = [
    {
        "name": "scene_01_hook_error",
        "html": "scene_01_hook_error.html",
        "duration": 8.0,
        "fps": 15,
    },
    {
        "name": "scene_02_data_gate",
        "html": "scene_02_data_gate.html",
        "duration": 9.0,
        "fps": 15,
    },
    {
        "name": "scene_03_download_success",
        "html": "scene_03_download_success.html",
        "duration": 10.0,
        "fps": 15,
    },
]

WIDTH  = 1080
HEIGHT = 1920


def ensure_dirs():
    """确保所有输出目录存在"""
    for d in [FRAMES_DIR, CLIPS_DIR, CONTACTS_DIR]:
        os.makedirs(d, exist_ok=True)


def render_scene(scene, playwright_ctx):
    """
    使用 Playwright 渲染单个 scene 的帧序列。

    策略: 页面加载后 CSS 动画自动播放，按时间戳截图。
    """
    name   = scene["name"]
    html   = scene["html"]
    dur    = scene["duration"]
    fps    = scene["fps"]
    total  = int(dur * fps)
    interval = 1.0 / fps

    html_path = os.path.join(HTML_DIR, html)
    if not os.path.exists(html_path):
        print(f"  [错误] HTML 文件不存在: {html_path}")
        return False

    out_dir = os.path.join(FRAMES_DIR, name)
    os.makedirs(out_dir, exist_ok=True)

    # 清空旧帧
    existing = [f for f in os.listdir(out_dir) if f.endswith(".png")]
    if existing:
        print(f"  [清除] {len(existing)} 个旧帧")
        for f in existing:
            os.remove(os.path.join(out_dir, f))

    browser = playwright_ctx.chromium.launch(headless=True)
    page = browser.new_page(
        viewport={"width": WIDTH, "height": HEIGHT},
        device_scale_factor=1,
    )

    # OOM 防护：关闭不需要的功能
    page.set_default_timeout(30000)

    file_url = "file:///" + html_path.replace("\\", "/")
    print(f"  [打开] {file_url}")

    try:
        page.goto(file_url, wait_until="networkidle", timeout=20000)
    except Exception as e:
        print(f"  [警告] goto 超时但继续: {e}")

    # 给 CSS 动画一个稳定的初始状态
    time.sleep(0.3)

    print(f"  [渲染] {total} 帧 ({dur}s @ {fps}fps)")
    start_time = time.time()
    captured = 0

    for i in range(total):
        # 计算目标时间戳
        target_ts = i * interval

        # 等待到目标时间
        while True:
            elapsed = time.time() - start_time
            if elapsed >= target_ts:
                break
            # 睡眠剩余时间的 80%，避免忙等
            remaining = target_ts - elapsed
            if remaining > 0.005:
                time.sleep(remaining * 0.8)
            else:
                break

        frame_path = os.path.join(out_dir, f"frame_{i:05d}.png")
        try:
            page.screenshot(path=frame_path, full_page=False)
            captured += 1
        except Exception as e:
            print(f"  [错误] 截图失败 frame {i}: {e}")

        # 进度报告
        if (i + 1) % 30 == 0 or i == total - 1:
            pct = (i + 1) / total * 100
            print(f"  [进度] {i+1}/{total} ({pct:.0f}%)")

    elapsed = time.time() - start_time
    print(f"  [完成] {captured}/{total} 帧, 实际耗时 {elapsed:.1f}s "
          f"(速率 {captured/elapsed:.1f} fps)")

    browser.close()
    return captured == total


def render_all_scenes():
    """依次渲染所有 scene"""
    from playwright.sync_api import sync_playwright

    print("=" * 60)
    print("Phase 3C — Scene Engine 帧序列渲染")
    print("=" * 60)

    results = {}
    with sync_playwright() as p:
        for scene in SCENES:
            name = scene["name"]
            print(f"\n▶ {name}")
            ok = render_scene(scene, p)
            results[name] = ok

    print("\n" + "=" * 60)
    print("渲染结果汇总:")
    for name, ok in results.items():
        status = "✅" if ok else "❌"
        print(f"  {status} {name}")
    print("=" * 60)
    return all(results.values())


def compose_clips():
    """用 FFmpeg 将 PNG 序列合成为 mp4 clip"""
    print("\n" + "=" * 60)
    print("合成 Scene Clips (FFmpeg)")
    print("=" * 60)

    all_ok = True
    for scene in SCENES:
        name     = scene["name"]
        fps      = scene["fps"]
        frames_dir = os.path.join(FRAMES_DIR, name)
        output    = os.path.join(CLIPS_DIR, f"{name}.mp4")

        if not os.path.isdir(frames_dir):
            print(f"  [跳过] 帧目录不存在: {frames_dir}")
            all_ok = False
            continue

        # 检查是否有帧文件
        png_files = [f for f in os.listdir(frames_dir) if f.endswith(".png")]
        if not png_files:
            print(f"  [跳过] 无帧文件: {frames_dir}")
            all_ok = False
            continue

        # FFmpeg: PNG 序列 → H.264 mp4
        input_pattern = os.path.join(frames_dir, "frame_%05d.png").replace("\\", "/")

        cmd = [
            "ffmpeg", "-y",
            "-framerate", str(fps),
            "-i", input_pattern,
            "-c:v", "libx264",
            "-pix_fmt", "yuv420p",
            "-preset", "medium",
            "-crf", "18",
            "-vf", f"scale={WIDTH}:{HEIGHT}:flags=lanczos",
            "-an",
            output,
        ]

        print(f"\n▶ {name}")
        print(f"   输入: {input_pattern}")
        print(f"   输出: {output}")

        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=300,
            )
            if result.returncode == 0:
                file_size = os.path.getsize(output) / (1024 * 1024)
                print(f"   ✅ 合成成功 ({file_size:.1f} MB)")
            else:
                print(f"   ❌ FFmpeg 错误:")
                print(result.stderr[-500:] if result.stderr else "无错误信息")
                all_ok = False
        except subprocess.TimeoutExpired:
            print(f"   ❌ FFmpeg 超时")
            all_ok = False
        except FileNotFoundError:
            print(f"   ❌ 未找到 FFmpeg (请确认 ffmpeg 在 PATH 中)")
            all_ok = False

    return all_ok


def generate_contact_sheets():
    """为每个 scene 生成 contact sheet (关键帧拼贴)"""
    print("\n" + "=" * 60)
    print("生成 Contact Sheets (FFmpeg tile)")
    print("=" * 60)

    all_ok = True
    for scene in SCENES:
        name     = scene["name"]
        frames_dir = os.path.join(FRAMES_DIR, name)
        output    = os.path.join(CONTACTS_DIR, f"{name}_contact.jpg")

        if not os.path.isdir(frames_dir):
            all_ok = False
            continue

        input_pattern = os.path.join(frames_dir, "frame_%05d.png").replace("\\", "/")

        # 6×6 = 36 缩略图，均匀采样
        cmd = [
            "ffmpeg", "-y",
            "-i", input_pattern,
            "-vf", "select='not(mod(n,10))',scale=160:284,tile=6x6",
            "-frames:v", "1",
            "-q:v", "3",
            output,
        ]

        print(f"\n▶ {name}")
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
            if result.returncode == 0:
                file_size = os.path.getsize(output) / 1024
                print(f"   ✅ Contact sheet 生成 ({file_size:.0f} KB)")
            else:
                print(f"   ❌ FFmpeg 错误: {result.stderr[-300:]}")
                all_ok = False
        except Exception as e:
            print(f"   ❌ 错误: {e}")
            all_ok = False

    return all_ok


def verify_outputs():
    """验证所有输出文件是否存在且符合预期"""
    print("\n" + "=" * 60)
    print("输出验证")
    print("=" * 60)

    all_ok = True
    checks = []

    for scene in SCENES:
        name = scene["name"]

        # 帧序列
        frames_dir = os.path.join(FRAMES_DIR, name)
        expected_frames = int(scene["duration"] * scene["fps"])
        if os.path.isdir(frames_dir):
            actual = len([f for f in os.listdir(frames_dir) if f.endswith(".png")])
            ok = actual >= expected_frames * 0.95  # 允许 5% 丢帧
            checks.append((f"{name} 帧序列", "✅" if ok else "❌",
                          f"{actual}/{expected_frames}"))
            if not ok:
                all_ok = False
        else:
            checks.append((f"{name} 帧序列", "❌", "目录不存在"))
            all_ok = False

        # Clip
        clip_path = os.path.join(CLIPS_DIR, f"{name}.mp4")
        if os.path.exists(clip_path):
            size_mb = os.path.getsize(clip_path) / (1024 * 1024)
            checks.append((f"{name} clip", "✅", f"{size_mb:.1f} MB"))
        else:
            checks.append((f"{name} clip", "❌", "文件不存在"))
            all_ok = False

        # Contact sheet
        cs_path = os.path.join(CONTACTS_DIR, f"{name}_contact.jpg")
        if os.path.exists(cs_path):
            size_kb = os.path.getsize(cs_path) / 1024
            checks.append((f"{name} contact sheet", "✅", f"{size_kb:.0f} KB"))
        else:
            checks.append((f"{name} contact sheet", "❌", "文件不存在"))
            all_ok = False

    print(f"\n{'项目':<30} {'状态':<6} {'详情':<15}")
    print("-" * 51)
    for item, status, detail in checks:
        print(f"{item:<30} {status:<6} {detail:<15}")

    return all_ok


def main():
    ensure_dirs()

    # Step 1: Render frame sequences
    print("\n═══ 第一步: Playwright 渲染帧序列 ═══")
    render_ok = render_all_scenes()

    # Step 2: Compose clips
    print("\n═══ 第二步: FFmpeg 合成 ═══")
    clip_ok = compose_clips()

    # Step 3: Contact sheets
    print("\n═══ 第三步: Contact Sheet ═══")
    cs_ok = generate_contact_sheets()

    # Step 4: Verify
    print("\n═══ 第四步: 输出验证 ═══")
    verify_ok = verify_outputs()

    all_ok = render_ok and clip_ok and cs_ok and verify_ok

    print("\n" + "=" * 60)
    if all_ok:
        print("🎉 Phase 3C 全部渲染完成!")
    else:
        print("⚠️  部分步骤未完成，请检查以上日志")
    print("=" * 60)

    return 0 if all_ok else 1


if __name__ == "__main__":
    sys.exit(main())